module time_schemes_mod
  use gptl_timer_mod
  use spark_cubed_sphere
  use state_mod
  use diag_mod
  use tend_mod
  use static_mod
  use params_mod
  use operators_adv_mod
  use operators_sw_mod
  use operators_nh_mod
  use ghost_interp_mod
  use hope_recon_mod
  use math_mod
  use jab_mod

  implicit none
  
  private kl, ku
    
  real(r_kind), dimension(:,:), allocatable :: IRK_A
  real(r_kind), dimension(:  ), allocatable :: IRK_b
  real(r_kind), dimension(:  ), allocatable :: IRK_c
  
  real(r_kind), dimension(:,:), allocatable :: ERK_A
  real(r_kind), dimension(:  ), allocatable :: ERK_b
  real(r_kind), dimension(:  ), allocatable :: ERK_c
  
  real(r_kind), dimension(:,:), allocatable :: Imatrix
  
  integer :: nElement
  
  integer, parameter :: kl=23, ku=23

  interface
    subroutine operators_run_interface(state, tend, substep)
      import state_type, diag_type, tend_type
      type(state_type), intent(inout) :: state
      type(tend_type ), intent(inout) :: tend
      integer, intent(in) :: substep
    end subroutine operators_run_interface
    subroutine time_integrator_interface(old, new, dt, state, tend)
      import state_type, tend_type, r_kind
      integer, intent(in) :: old
      integer, intent(in) :: new
      real(r_kind), intent(in) :: dt
      type(state_type), intent(inout) :: state(:)
      type(tend_type ), intent(inout) :: tend (:)
    end subroutine time_integrator_interface
  end interface

  procedure(operators_run_interface), pointer :: operators_run => null()

contains

  subroutine time_schemes_init()
    
    select case(model_type)
      case(0)
        operators_run => operators_adv_run
      case(1)
        operators_run => operators_sw_run
      case(2)
        operators_run => operators_nh_run
      end select
      
      if( trim(time_scheme) == 'IMEX' )then
        call init_IMEX
      endif
      
  end subroutine time_schemes_init

  subroutine euler(old, new, dt, state, tend)

    integer, intent(in) :: old
    integer, intent(in) :: new
    real(r_kind), intent(in) :: dt
    type(state_type), intent(inout) :: state(:)
    type(tend_type ), intent(inout) :: tend (:)
    
    if(use_tci)then
      call find_trouble_cells(state(old), diag)
    endif

    call operators_run(state(old), tend(new), substep=1)
    call update_state(dt, tend(new), state(old), state(new))

  end subroutine euler

  subroutine ssp_rk3(old, new, dt, state, tend)

    integer, intent(in) :: old
    integer, intent(in) :: new
    real(r_kind), intent(in) :: dt
    type(state_type), intent(inout) :: state(:)
    type(tend_type ), intent(inout) :: tend (:)

    integer k1, k2, k3
    integer idom, ids, ide, jds, jde, kds, kde, i, j, k, iVar
    real(r_kind), parameter :: c1 = 0.25_r8, c2 = 2.0_r8 / 3.0_r8

    k1 = old; k2 = 3; k3 = new
    
    if(use_tci)then
      call find_trouble_cells(state(old), diag)
    endif
    
    call operators_run(state(k1), tend(k1), substep=1)
    call update_state(dt, tend(k1), state(old), state(k2))

    call operators_run(state(k2), tend(k2), substep=2)
    
    do idom = 1, mesh%ndom
      if (state(old)%mesh%domain(idom)%initialized) then
        call state(old)%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                tend(k2)%domain(idom)%q(i,j,k,iVar) = c1 * ( tend(k1)%domain(idom)%q(i,j,k,iVar) &
                                                           + tend(k2)%domain(idom)%q(i,j,k,iVar) )
              end do
            end do
          end do
        enddo
        !$OMP END PARALLEL DO
      end if
    end do
    call update_state(dt, tend(k2), state(old), state(k3))

    call operators_run(state(k3), tend(k3), substep=3)
    
    do idom = 1, mesh%ndom
      if (state(old)%mesh%domain(idom)%initialized) then
        call state(old)%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                tend(k3)%domain(idom)%q(i,j,k,iVar) = c2 * ( tend(k2)%domain(idom)%q(i,j,k,iVar) &
                                                           + tend(k3)%domain(idom)%q(i,j,k,iVar) )
              end do
            end do
          end do
        enddo
        !$OMP END PARALLEL DO
      end if
    end do
    call update_state(dt, tend(k3), state(old), state(new))

  end subroutine ssp_rk3

  subroutine WRF_RK3(old, new, dt, state, tend)

    integer, intent(in) :: old
    integer, intent(in) :: new
    real(r_kind), intent(in) :: dt
    type(state_type), intent(inout) :: state(:)
    type(tend_type ), intent(inout) :: tend (:)

    real(r_kind), parameter :: w(3) = [1./3.,0.5,1.]
    integer k1, k2, k3
    integer idom, ids, ide, jds, jde, kds, kde, i, j, k, iVar

    k1 = old; k2 = 3; k3 = new
    
    if(use_tci)then
      call find_trouble_cells(state(old), diag)
    endif
    
    call operators_run(state(k1), tend(k1), substep=1)
    call update_state(w(1) * dt, tend(k1), state(old), state(k2))

    call operators_run(state(k2), tend(k2), substep=2)
    call update_state(w(2) * dt, tend(k2), state(old), state(k3))

    call operators_run(state(k3), tend(k3), substep=3)
    call update_state(w(3) * dt, tend(k3), state(old), state(new))

  end subroutine WRF_RK3

  subroutine rk4(old, new, dt, state, tend)

    integer, intent(in) :: old
    integer, intent(in) :: new
    real(r_kind), intent(in) :: dt
    type(state_type), intent(inout) :: state(:)
    type(tend_type ), intent(inout) :: tend (:)

    real(r_kind), parameter :: w(4) = [1./6.,1./3.,1./3.,1./6.]
    integer k1, k2, k3, k4
    integer idom, ids, ide, jds, jde, kds, kde, i, j, k, iVar

    k1 = old; k2 = 3; k3 = 4; k4 = new
    
    if(use_tci)then
      call find_trouble_cells(state(old), diag)
    endif

    call operators_run(state(k1), tend(k1), substep=1)
    call update_state(0.5_r8 * dt, tend(k1), state(old), state(k2))

    call operators_run(state(k2), tend(k2), substep=2)
    call update_state(0.5_r8 * dt, tend(k2), state(old), state(k3))

    call operators_run(state(k3), tend(k3), substep=3)
    call update_state(         dt, tend(k3), state(old), state(k4))

    call operators_run(state(k4), tend(k4), substep=4)
    
    do idom = 1, mesh%ndom
      if (state(old)%mesh%domain(idom)%active) then
        call state(old)%mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1, nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                tend(new)%domain(idom)%q(i,j,k,iVar) = w(1) * tend(k1)%domain(idom)%q(i,j,k,iVar) &
                                                     + w(2) * tend(k2)%domain(idom)%q(i,j,k,iVar) &
                                                     + w(3) * tend(k3)%domain(idom)%q(i,j,k,iVar) &
                                                     + w(4) * tend(k4)%domain(idom)%q(i,j,k,iVar)
              end do
            end do
          end do
        enddo
        !$OMP END PARALLEL DO
      end if
    end do
    call update_state(dt, tend(new), state(old), state(new))

  end subroutine rk4
 
  subroutine update_state(dt, tend, old_state, new_state)

    real(r_kind), intent(in) :: dt
    type(tend_type ), intent(in   ) :: tend
    type(state_type), intent(in   ) :: old_state
    type(state_type), intent(inout) :: new_state

    integer idom, ids, ide, jds, jde, kds, kde, i, j, k, iVar
    
    ret = gptlstart('update_state')

    associate (mesh => old_state%mesh)
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        !$OMP PARALLEL DO COLLAPSE(4)
        do iVar = 1,nVar
          do k = kds, kde
            do j = jds, jde
              do i = ids, ide
                new_state%domain(idom)%q(i,j,k,iVar) = old_state%domain(idom)%q(i,j,k,iVar) + dt * tend%domain(idom)%q(i,j,k,iVar)
              end do
            end do
          end do
        enddo
        !$OMP END PARALLEL DO
      end if
    end do
    end associate

    call fill_halo(new_state%array)
    call ghost_interp(new_state, diag)
    
    ret = gptlstop('update_state')

  end subroutine update_state
  
  subroutine init_IMEX
    real(r_kind) :: alpha
    real(r_kind) :: beta
    real(r_kind) :: lambda
    real(r_kind) :: eta
    real(r_kind) :: sigma
    real(r_kind) :: gamma
    
    integer idom, ids, ide, jds, jde, kds, kde
    integer i, j, k, iVar
    
    nElement = nVar * nz
  
    ! For matrices
    allocate(Imatrix(nElement,nElement))
    
    Imatrix = 0
    
    do i = 1,nElement
      Imatrix(i,i) = 1
    enddo
    
    ! For IRK coef
    allocate(IRK_A(IRK_stage+1,IRK_stage+1))
    allocate(IRK_b(IRK_stage+1))
    allocate(IRK_c(IRK_stage+1))
    
    ! For ERK coef
    allocate(ERK_A(IRK_stage+1,IRK_stage+1))
    allocate(ERK_b(IRK_stage+1))
    allocate(ERK_c(IRK_stage+1))
    
    select case(IRK_stage)
    case(2)
      ! Ascher, 1997, section 2.4
      gamma = ( 3. + sqrt(3.) ) / 6.
      IRK_A = 0
      IRK_A(2,2) = gamma
      IRK_A(3,2) = 1. - 2. * gamma; IRK_A(3,3) = gamma
      IRK_c = [0._r_kind,gamma,1.-gamma]
      IRK_b = [0._r_kind,0.5_r_kind,0.5_r_kind]
      
      ERK_A = 0
      ERK_A(2,1) = gamma
      ERK_A(3,1) = gamma - 1.; ERK_A(3,2) = 2. * ( 1. - gamma )
      ERK_c = [0._r_kind, gamma, 1. - gamma]
      ERK_b = [0._r_kind,0.5_r_kind,0.5_r_kind]
    case(3)
      ! Ascher, 1997, section 2.7
      lambda = 0.435866521508459
      IRK_A = 0
      IRK_A(2,2) = lambda
      IRK_A(3,2) = (1-lambda)/2;                 IRK_A(3,3) = lambda
      IRK_A(4,2) = (-6*lambda**2+16*lambda-1)/4; IRK_A(4,3) = (6*lambda**2-20*lambda+5)/4; IRK_A(4,4) = lambda
      IRK_c = [0._r_kind,lambda,(1._r_kind+lambda)/2._r_kind,1._r_kind]
      IRK_b = IRK_A(4,:)
      
      ERK_A = 0
      ERK_A(2,1) = lambda
      ERK_A(3,1) = 0.3212788860; ERK_A(3,2) = 0.3966543747
      ERK_A(4,1) =-0.105858296 ; ERK_A(4,2) = 0.5529291479; ERK_A(4,3) = 0.5529291479
      ERK_c = IRK_c
      ERK_b = IRK_b
      
    case default
      print*,'Unknown IRK_stage in init_IMEX'
      stop
    end select
    
  end subroutine init_IMEX
  
  subroutine IMEX(old, new, dt, state, tend)
    integer, intent(in) :: old
    integer, intent(in) :: new
    real(r_kind), intent(in) :: dt
    type(state_type), intent(inout) :: state(:)
    type(tend_type ), intent(inout) :: tend (:)
    
    integer :: iStage, jStage, kStage
    integer :: k1, k2, k3, k4
    integer :: ims, ime, jms, jme, kms, kme
    integer :: ids, ide, jds, jde, kds, kde
    integer :: idom, i, j, k, iVar
    integer :: niter

    k1 = 1; k2 = 2; k3 = 3; k4 = 4
    
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, &
                                        ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme )
        
        state_IMEX(k1)%domain(idom)%q = state(old)%domain(idom)%q
        state_IM      %domain(idom)%q = state(old)%domain(idom)%q
      endif
    enddo
    
    do iStage = 1, IRK_stage+1
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%active) then
          call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde, &
                                          ims=ims, ime=ime, jms=jms, jme=jme, kms=kms, kme=kme )
                                          
          tend_IM(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = 0
          tend_EX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = 0
          if( iStage > 1 )then
            state_IMEX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = 0
            do jStage = 1, iStage-1
              state_IMEX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:)                              &
              = state_IMEX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:)                            &
              + dt * ( IRK_A(iStage,jStage) * tend_IM(jStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
                     + ERK_A(iStage,jStage) * tend_EX(jStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) )
            enddo
            state_IMEX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = state(old)        %domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
                                                                         + state_IMEX(iStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:)
          endif ! iStage > 1
        endif ! mesh%domain(idom)%active
      enddo ! idom
      
      if(IRK_A(iStage,iStage)/=0) call IRK_z(tend_IM(iStage),state_IM,state_IMEX(iStage),IRK_A(iStage,iStage)*dt,iStage)
      !if(IRK_A(iStage,iStage)/=0) call RKR(tend_IM(iStage),state_IM,state_IMEX(iStage),IRK_A(iStage,iStage)*dt)
      
      call fill_halo(state_IM%array)
      call ghost_interp(state_IM, diag)
      
      call operators_hor_run(state_IM, tend_EX(iStage))
    enddo
    
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        state(new)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = 0
        do jStage = 1, IRK_stage+1
          state(new)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
          = state(new)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
          + dt * ( IRK_b(jStage) * tend_IM(jStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
                 + ERK_b(jStage) * tend_EX(jStage)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) )
        enddo
        
        state(new)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) = state(old)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:) &
                                                             + state(new)%domain(idom)%q(ids:ide,jds:jde,kds:kde,:)
      endif ! mesh%domain(idom)%active
    enddo
    
    call Rayleigh_damping(state(new),state_ref,dvs,dve)

    call fill_halo(state(new)%array)
    call ghost_interp(state(new), diag)
    
  end subroutine IMEX
  
  subroutine IRK_z(tend,state_new,state_old,dt,iStage)
    type(tend_type ), intent(inout) :: tend
    type(state_type), intent(inout) :: state_new
    type(state_type), intent(in   ) :: state_old
    real(r_kind)    , intent(in   ) :: dt
    integer(i_kind) , intent(in   ) :: iStage
    
    real(r_kind), dimension(nElement,nElement) :: Amatrix
    real(r_kind), dimension(nElement         ) :: Rmatrix
    real(r_kind), dimension(nElement,1       ) :: dq
    
    integer      :: iter,jter,kter
    real(r_kind) :: residual1,residual2
    real(r_kind) :: dresidual ! relative error
    integer      :: newton_scheme_num = 1
    integer      :: max_sub_iter = 30
    integer      :: max_scheme_num = 3
    
    integer      :: ids, ide, jds, jde, kds, kde
    integer      :: idom,i,j,k
    real(r_kind) :: lambda
    logical      :: reset_loop
    
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        associate(niter => diag%domain(idom)%niter)
          call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
          !$OMP PARALLEL DO PRIVATE(i,residual1,residual2,dresidual,iter,lambda,Amatrix,dq,Rmatrix,newton_scheme_num) COLLAPSE(2)
          do j = jds,jde
            do i = ids,ide
              lambda = 1
              reset_loop = .true.
              iter = 0
              jter = 0 ! sub loop
              newton_scheme_num = 1
              do while(reset_loop)
                state_IRK%domain(idom)%q(i,j,kds:kde,:) = state_old%domain(idom)%q(i,j,kds:kde,:)
                
                residual1 = 1.e15
                residual2 = 1.e15
                dresidual = 1.e15
                dq        = 0
                  
                do while( residual2>IRK_residual .and. iter<IRK_max_iter .and. dresidual> 1.e-10 )
                !do while( residual2>IRK_residual .and. iter<IRK_max_iter )
                  !select case(newton_scheme_num)
                  !case(1)
                  !  call newton_scheme_1(state_IRK,state_old,Amatrix,dq,Rmatrix,idom,i,j,kds,kde)
                  !case(2)
                  !  call newton_scheme_2(state_IRK,state_old,Amatrix,dq,Rmatrix,idom,i,j,kds,kde)
                  !case(3)
                  !  call newton_scheme_3(state_IRK,state_old,Amatrix,dq,Rmatrix,idom,i,j,kds,kde)
                  !case(4)
                    call newton_scheme_4(state_IRK,state_old,Amatrix,dq,Rmatrix,idom,i,j,kds,kde)
                  !end select
                  
                  ! Update stat
                  state_IRK%domain(idom)%q(i,j,kds:kde,:) = state_IRK%domain(idom)%q(i,j,kds:kde,:) + lambda * transpose(reshape(dq,(/nVar,nz/)))
                  
                  residual1 = residual2
                  residual2 = norm2(dq)
                  !residual2 = norm2( matmul( transpose(Amatrix), Rmatrix ) )
                  dresidual = abs(residual2-residual1)/abs(residual1)
                  iter      = iter + 1
                  jter      = jter + 1
                  
                  !print*,'i,j,iter,residual2 ',i,j,iter,residual2
                  
                  !open(123,file='A.txt')
                  !do k = 1,240
                  !  write(123,'(20000f50.15)')Amatrix(k,:)
                  !enddo
                  !close(123)
                  !open(123,file='R.txt')
                  !do k = 1,240
                  !  write(123,'(20000f50.15)')Rmatrix(k)
                  !enddo
                  !close(123)
                  !stop 'Check IRK_z'
                  
                  reset_loop = .false.
                  if( isnan(residual2) )then
                    print*,'Nan appears during implicit iteration, idom,i,j,iter',idom,i,j,iter
                    stop
                  endif
                  
                  !if( isnan(residual2) )then
                  !  if(newton_scheme_num<max_scheme_num)then
                  !    newton_scheme_num = newton_scheme_num + 1
                  !  else
                  !    newton_scheme_num = 1
                  !    lambda = 0.125 * lambda
                  !  endif
                  !  jter = 0
                  !  reset_loop = .true.
                  !elseif( mod(iter,max_sub_iter)==0 .and. iter>=max_sub_iter )then
                  !  if(newton_scheme_num<max_scheme_num)then
                  !    newton_scheme_num = newton_scheme_num + 1
                  !  else
                  !    newton_scheme_num = 1
                  !    lambda = 0.85**(iter/100.) * lambda
                  !  endif
                  !  jter = 0
                  !  reset_loop = .false.
                  !!elseif( jter>20 .and. newton_scheme_num==max_scheme_num )then
                  !!  lambda = 0.9 * lambda
                  !!  reset_loop = .true.
                  !else
                  !  reset_loop = .false.
                  !endif
                enddo
              enddo
              
              niter(i,j) = iter
              
              call operators_ver_run(state_IRK, tend_IRK, idom, i, j) 
              
              state_new%domain(idom)%q(i,j,kds:kde,:) = state_old%domain(idom)%q(i,j,kds:kde,:) + dt * tend_IRK%domain(idom)%q(i,j,kds:kde,:)
              
              tend%domain(idom)%q(i,j,kds:kde,:) = tend_IRK%domain(idom)%q(i,j,kds:kde,:)
              
            enddo
          enddo
          !$OMP END PARALLEL DO
          print*,'IMEX Stage',iStage,'/',IRK_stage+1,' done on idom',idom,' with min/max/mean iter', minval(niter), maxval(niter), sum(float(niter)) / float( size(niter,1) * size(niter,2) )
        end associate
      endif
    enddo
    
  end subroutine IRK_z
  
  subroutine RKR(tend,state_new,state_old,dt)
    type(tend_type ), intent(inout) :: tend
    type(state_type), intent(inout) :: state_new
    type(state_type), intent(inout) :: state_old
    real(r_kind)    , intent(in   ) :: dt
    
    
    real(r_kind), dimension(nElement,nElement) :: Amatrix
    real(r_kind), dimension(nElement,nElement) :: Jmatrix
    real(r_kind), dimension(nElement         ) :: Rmatrix
    real(r_kind), dimension(nElement         ) :: r
    real(r_kind), dimension(nElement         ) :: dq
    
    integer :: ids, ide, jds, jde, kds, kde
    integer :: idom,i,j,k,iVar
    
    do idom = 1, mesh%ndom
      if (mesh%domain(idom)%active) then
        call mesh%get_params(idom=idom, ids=ids, ide=ide, jds=jds, jde=jde, kds=kds, kde=kde)
        !$OMP PARALLEL DO PRIVATE(Jmatrix,Amatrix,r,dq) COLLAPSE(2)
        do j = jds, jde
          do i = ids, ide
            call operators_ver_run(state_old, tend_HON(1), idom, i, j)
            
            call calc_flux_jab(Jmatrix,state_old,diag,idom,i,j,kds,kde)
            
            Amatrix = Imatrix / dt - Jmatrix
            
            r = reshape( transpose( tend_HON(1)%domain(idom)%q(i,j,kds:kde,:) ), (/nElement/) )
            
            ! Solve A*dq=R by Steepest Descent Method
            call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i,jdx=j)
            
            tend     %domain(idom)%q(i,j,kds:kde,:) = transpose( reshape(dq,(/nVar,nz/)) ) / dt
            state_new%domain(idom)%q(i,j,kds:kde,:) = state_old%domain(idom)%q(i,j,kds:kde,:) + dt * tend%domain(idom)%q(i,j,kds:kde,:)
          enddo
        enddo
        !$OMP END PARALLEL DO
      endif
    enddo
    
  end subroutine RKR
      
  ! Original Newton method, 2nd order
  subroutine newton_scheme_1(state,state_old,Amatrix,dq,r,idom,i,j,kds,kde)
    type(state_type), intent(inout) :: state ! input only
    type(state_type), intent(in   ) :: state_old
    integer(i_kind),  intent(in   ) :: idom, i, j
    integer(i_kind),  intent(in   ) :: kds, kde
    real   (r_kind), dimension(nElement,nElement), intent(out) :: Amatrix
    real   (r_kind), dimension(nElement         ), intent(out) :: dq
    real   (r_kind), dimension(nElement         ), intent(out) :: r
    
    real(r_kind), dimension(nElement,nElement) :: Jmatrix
    
    ! Stage 1
    call operators_ver_run(state, tend_HON(1), idom, i, j)
    
    call calc_flux_jab(Jmatrix,state,diag,idom,i,j,kds,kde)
    
    Amatrix = Imatrix / dt - Jmatrix
    r       = calc_residual(state,state_old,tend_HON(1),dt,idom,i,j,kds,kde)
    
    ! Solve A*dq=R by Steepest Descent Method
    call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i,jdx=j)
    
  end subroutine newton_scheme_1
      
  ! Xiaowu Li, Zhinan Wu, Lin Wang and Qian Zhang, 2013, 9th order
  subroutine newton_scheme_4(state,state_old,Amatrix,dq,r,idom,i,j,kds,kde)
    type(state_type), intent(inout) :: state ! input only
    type(state_type), intent(in   ) :: state_old
    integer(i_kind),  intent(in   ) :: idom, i, j
    integer(i_kind),  intent(in   ) :: kds, kde
    real   (r_kind), dimension(nElement,nElement), intent(out) :: Amatrix
    real   (r_kind), dimension(nElement         ), intent(out) :: dq
    real   (r_kind), dimension(nElement         ), intent(out) :: r
    
    real(r_kind), dimension(nElement,nElement) :: Jmatrix
    
    integer :: k1, k2, k3, k4
    
    k1 = 1; k2 = 2; k3 = 3; k4 = 4
        
    ! Stage 1
    call operators_ver_run(state, tend_HON(k1), idom, i, j)
    r = calc_residual(state,state_old,tend_HON(k1),dt,idom,i,j,kds,kde)
    call calc_flux_jab(Jmatrix,state,diag,idom,i,j,kds,kde)
    Amatrix = Imatrix / dt - Jmatrix
    call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i)
    state_HON(k1)%domain(idom)%q(i,j,kds:kde,:) = state%domain(idom)%q(i,j,kds:kde,:) + transpose(reshape(dq,(/nVar,nz/))) ! y_k
    
    ! Stage 2
    call operators_ver_run(state_HON(k1), tend_HON(k2), idom, i, j)
    r = calc_residual(state_HON(k1),state_old,tend_HON(k2),dt,idom,i,j,kds,kde)
    call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i,jdx=j)
    state_HON(k2)%domain(idom)%q(i,j,kds:kde,:) = state_HON(k1)%domain(idom)%q(i,j,kds:kde,:) + transpose(reshape(dq,(/nVar,nz/))) ! z_k
    
    ! Stage 3
    call operators_ver_run(state_HON(k2), tend_HON(k3), idom, i, j)
    r = calc_residual(state_HON(k2),state_old,tend_HON(k3),dt,idom,i,j,kds,kde)
    call calc_flux_jab(Jmatrix,state_HON(k2),diag,idom,i,j,kds,kde)
    Amatrix = Imatrix / dt - Jmatrix
    call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i,jdx=j)
    state_HON(k3)%domain(idom)%q(i,j,kds:kde,:) = state_HON(k2)%domain(idom)%q(i,j,kds:kde,:) + transpose(reshape(dq,(/nVar,nz/))) ! p_k
    
    ! Stage 4
    call operators_ver_run(state_HON(k3), tend_HON(k4), idom, i, j)
    r = calc_residual(state_HON(k3),state_old,tend_HON(k4),dt,idom,i,j,kds,kde)
    call lapack_dgbsv(Amatrix,r,dq,nElement,kl=kl,ku=ku,idx=i,jdx=j)
    
    state%domain(idom)%q(i,j,kds:kde,:) = state_HON(k3)%domain(idom)%q(i,j,kds:kde,:)
    
  end subroutine newton_scheme_4
  
  function calc_residual(state,state_old,tend,dt,idom,i,j,kds,kde) result (r)
    real(r_kind), dimension(nElement) :: r
    type(state_type), intent(inout) :: state
    type(state_type), intent(in   ) :: state_old
    type(tend_type ), intent(inout) :: tend
    real   (r_kind) , intent(in   ) :: dt
    integer(i_kind) , intent(in   ) :: idom
    integer(i_kind) , intent(in   ) :: i, j
    integer(i_kind) , intent(in   ) :: kds,kde
  
    real(r_kind), dimension(nVar,kds:kde) :: r_2d
    
    r_2d = transpose( -( state%domain(idom)%q(i,j,kds:kde,:) - state_old%domain(idom)%q(i,j,kds:kde,:) ) / dt &
                      + tend%domain(idom)%q(i,j,kds:kde,:) )
    r = reshape(r_2d,(/nElement/))
  end function calc_residual
  
  subroutine calc_flux_jab(Jmatrix,state,diag,idom,i,j,kds,kde)
    real   (r_kind)    , intent(out) :: Jmatrix(nElement,nElement)
    integer(i_kind)    , intent(in ) :: idom,i,j
    type   (state_type), intent(in ) :: state
    type   (diag_type ), intent(in ) :: diag
    integer(i_kind)    , intent(in ) :: kds
    integer(i_kind)    , intent(in ) :: kde
    
    real(r_kind), dimension(nVar,nElement,kds:kde+1) :: dFdq
    real(r_kind), dimension(nVar,nVar)               :: dFdqL,dFdqR
    real(r_kind), dimension(nVar,nElement)           :: dqRdq,dqLdq
    real(r_kind), dimension(nVar,nVar)               :: dpdq
      
    real(r_kind), dimension(nVar,nElement,nz+1) :: tangent_coefP
    real(r_kind), dimension(nVar,nElement,nz+1) :: tangent_coefN
    
    real(r_kind), dimension(swv) :: dqRdq_local, dqLdq_local
    
    real(r_kind) :: JhL, JhR
    real(r_kind), dimension(3,3) :: iG, iGL, iGR
    
    integer :: pc, pes(6), pee(6)
    integer :: irow, irows, irowe
    integer :: icol, icols, icole
    integer :: iVar,ic,jc,k,iElement
    integer :: vs, ve
    integer :: nrc, rwL, rwR
    integer :: ip, p, p_adj
    
    Jmatrix = 0
    dFdq    = 0
      
    associate (mesh => diag%mesh)
      call mesh%get_params(idom=idom, pc=pc, pes=pes, pee=pee)
      associate (Jab     => mesh %domain(idom)%Jab,      &
                 Jv      => mesh %domain(idom)%Jv,       &
                 iGv     => mesh %domain(idom)%iGv,      &
                 deta    => mesh %domain(idom)%deta,     &
                 gravity => static%domain(idom)%gravity, &
                 q       => state%domain(idom)%q,        &
                 qD      => diag%domain(idom)%qD,        &
                 qU      => diag%domain(idom)%qU,        &
                 recon_v => recon_arrays(idom)%recon_ver )
      ! Set reconstruction Jacobian
      !do k = kds, kde
      !  tangent_coefN(:,:,k+1) = 0
      !  tangent_coefP(:,:,k  ) = 0
      !  do iVar = 1,nVar
      !    nrc = recon_v(k)%ptr%nc
      !    rwL = recon_v(k)%ptr%rwL
      !    rwR = recon_v(k)%ptr%rwR
      !    call recon_v(k)%ptr%get_jab(dqLdq_local(1:nrc),dqRdq_local(1:nrc),q(i,j,k-rwL:k+rwR,iVar))
      !    
      !    do ic = 1, nrc
      !      iElement = ( k + ic - rwL - 2 ) * nVar + iVar
      !      tangent_coefN(iVar,iElement,k+1) = dqRdq_local(ic)
      !      tangent_coefP(iVar,iElement,k  ) = dqLdq_local(ic)
      !    enddo
      !  enddo
      !enddo
      !tangent_coefN(:,:,kds  ) = tangent_coefP(:,:,kds  )
      !tangent_coefP(:,:,kde+1) = tangent_coefN(:,:,kde+1)
      
      do k = kds, kde
        tangent_coefN(:,:,k+1) = 0
        tangent_coefP(:,:,k  ) = 0
        do iVar = 1,nVar
          nrc = recon_v(k)%ptr%nc
          rwL = recon_v(k)%ptr%rwL
          rwR = recon_v(k)%ptr%rwR
          call recon_v(k)%ptr%get_jab(dqLdq_local(1:nrc),dqRdq_local(1:nrc),q(i,j,k-rwL:k+rwR,iVar))
          
          do ic = 1, nrc
            iElement = ( k - rwL + ic - 2 ) * nVar + iVar
            tangent_coefN(iVar,iElement,k+1) = dqRdq_local(ic) * Jab(pes(6),i,j,k)
            tangent_coefP(iVar,iElement,k  ) = dqLdq_local(ic) * Jab(pes(5),i,j,k)
          enddo
        enddo
      enddo
      tangent_coefN(:,:,kds  ) = tangent_coefP(:,:,kds  )
      tangent_coefP(:,:,kde+1) = tangent_coefN(:,:,kde+1)
      
      p     = pes(5)
      p_adj = pes(6)
      ! Combine reconstruction Jacobian and Riemann solver Jacobian
      ! Bottom boundary
      k = kds
      dqLdq = tangent_coefN(:,:,k)
      iG = mesh%get_dom_iG(idom,p,i,j,k)
      call dpdq_mtx( dpdq, Jab(p,i,j,k), iG, qD (p,i,j,k,:) )
      dFdqL = dpdq
      dFdq(:,:,k) = matmul(dFdqL,dqLdq)
      
      ! Center region
      do k = kds+1, kde
        dqLdq = tangent_coefN(:,:,k)
        dqRdq = tangent_coefP(:,:,k)
        
        JhL = mesh%get_dom_Jh(idom,p_adj,i,j,k-1)
        JhR = mesh%get_dom_Jh(idom,p    ,i,j,k  )
        iGL = mesh%get_dom_iG(idom,p_adj,i,j,k-1)
        iGR = mesh%get_dom_iG(idom,p    ,i,j,k  )
        
        call AUSM_jab(dFdqL,dFdqR,Jab(    p_adj,i,j,k-1  ),Jab(    p,i,j,k  ),&
                                  JhL                     ,JhR               ,&
                                  Jv (    p_adj,i,j,k-1  ),Jv (    p,i,j,k  ),&
                                  iGL                     ,iGR               ,&
                                  iGv(:,:,p_adj,i,j,k-1  ),iGv(:,:,p,i,j,k  ),&
                                  qU (    p_adj,i,j,k-1,:),qD (    p,i,j,k,:) )
        
        dFdq(:,:,k) = matmul(dFdqL,dqLdq) + matmul(dFdqR,dqRdq)
      enddo
      
      ! Top boundary
      k = kde + 1
      dqLdq = tangent_coefN(:,:,k)
      iG = mesh%get_dom_iG(idom,p,i,j,k)
      call dpdq_mtx( dpdq, Jab(p,i,j,k), iG, qD (p,i,j,k,:) )
      dFdqL = dpdq
      dFdq(:,:,k) = matmul(dFdqL,dqLdq)
      
      ! Add source Jacobian
      do k = kds, kde
        irows = nVar * ( k - 1 ) + 1
        irowe = irows + nVar - 1
        
        ! Flux term
        Jmatrix(irows:irowe,:) = -( dFdq(:,:,k+1) - dFdq(:,:,k) ) / deta / Jab(pc,i,j,k)
        
        ! gravity (source) term
        Jmatrix(irows+3,(k-1)*nVar+1) = Jmatrix(irows+3,(k-1)*nVar+1) - gravity(i,j,k)
        Jmatrix(irows+3,(k-1)*nVar+6) = Jmatrix(irows+3,(k-1)*nVar+6) - gravity(i,j,k)
        
        !! Rayleigh damping term
        !if(case_num==2)then
        !  vs = 1
        !  ve = 6
        !elseif(case_num==5)then
        !  vs = 4
        !  ve = 4
        !endif
        !
        !do iVar = vs,ve
        !  irow = irows+iVar-1
        !  icol = irow
        !  Jmatrix(irow,icol) = Jmatrix(irow,icol) - relax_coef(i,j,k) / dt
        !enddo
      enddo
      
      end associate
    end associate
  end subroutine calc_flux_jab

end module time_schemes_mod
